#include <assert.h>
#include "avl.h"

/**
 * @brief 取较大值
 * @param lhs 左值
 * @param rhs 右值
 * @return 较大的值
 */
static uint32_t max(uint32_t lhs, uint32_t rhs) {
    return lhs < rhs ? rhs : lhs;
}

/**
 * @brief 维护节点的高度和计数字段
 * @param node AVL树节点
 */
static void avl_update(AVLNode *node) {
    node->height = 1 + max(avl_height(node->left), avl_height(node->right));
    node->cnt = 1 + avl_cnt(node->left) + avl_cnt(node->right);
}

/**
 * @brief 左旋转操作
 * @param node 要旋转的节点
 * @return 旋转后的新根节点
 */
static AVLNode *rot_left(AVLNode *node) {
    AVLNode *parent = node->parent;
    AVLNode *new_node = node->right;
    AVLNode *inner = new_node->left;
    
    // node <-> inner
    node->right = inner;
    if (inner) {
        inner->parent = node;
    }
    
    // parent <- new_node
    new_node->parent = parent;
    
    // new_node <-> node
    new_node->left = node;
    node->parent = new_node;
    
    // 更新辅助数据
    avl_update(node);
    avl_update(new_node);
    return new_node;
}

/**
 * @brief 右旋转操作
 * @param node 要旋转的节点
 * @return 旋转后的新根节点
 */
static AVLNode *rot_right(AVLNode *node) {
    AVLNode *parent = node->parent;
    AVLNode *new_node = node->left;
    AVLNode *inner = new_node->right;
    
    // node <-> inner
    node->left = inner;
    if (inner) {
        inner->parent = node;
    }
    
    // parent <- new_node
    new_node->parent = parent;
    
    // new_node <-> node
    new_node->right = node;
    node->parent = new_node;
    
    // 更新辅助数据
    avl_update(node);
    avl_update(new_node);
    return new_node;
}

/**
 * @brief 修复左子树高2的情况
 * @param node 需要修复的节点
 * @return 修复后的根节点
 */
static AVLNode *avl_fix_left(AVLNode *node) {
    if (avl_height(node->left->left) < avl_height(node->left->right)) {
        node->left = rot_left(node->left);  // 转换2
    }
    return rot_right(node);                 // 转换1
}

/**
 * @brief 修复右子树高2的情况
 * @param node 需要修复的节点
 * @return 修复后的根节点
 */
static AVLNode *avl_fix_right(AVLNode *node) {
    if (avl_height(node->right->right) < avl_height(node->right->left)) {
        node->right = rot_right(node->right);
    }
    return rot_left(node);
}

AVLNode *avl_fix(AVLNode *node) {
    while (true) {
        AVLNode **from = &node; // 保存修复后的子树
        AVLNode *parent = node->parent;
        if (parent) {
            // 将修复后的子树附加到父节点
            from = parent->left == node ? &parent->left : &parent->right;
        }   // 否则：保存到局部变量`node`
        
        // 更新辅助数据
        avl_update(node);
        
        // 修复高度差为2的情况
        uint32_t l = avl_height(node->left);
        uint32_t r = avl_height(node->right);
        if (l == r + 2) {
            *from = avl_fix_left(node);
        } else if (l + 2 == r) {
            *from = avl_fix_right(node);
        }
        
        // 根节点，停止
        if (!parent) {
            return *from;
        }
        
        // 继续到父节点，因为其高度可能已改变
        node = parent;
    }
}

/**
 * @brief 分离一个子节点为空的节点
 * @param node 要分离的节点
 * @return 分离后的新根节点
 */
static AVLNode *avl_del_easy(AVLNode *node) {
    assert(!node->left || !node->right);    // 最多1个子节点
    AVLNode *child = node->left ? node->left : node->right; // 可能为NULL
    AVLNode *parent = node->parent;
    
    // 更新子节点的父指针
    if (child) {
        child->parent = parent; // 可能为NULL
    }
    
    // 将子节点附加到祖父节点
    if (!parent) {
        return child;   // 移除根节点
    }
    AVLNode **from = parent->left == node ? &parent->left : &parent->right;
    *from = child;
    
    // 重新平衡更新后的树
    return avl_fix(parent);
}

AVLNode *avl_del(AVLNode *node) {
    // 0或1个子节点的简单情况
    if (!node->left || !node->right) {
        return avl_del_easy(node);
    }
    
    // 查找后继节点
    AVLNode *victim = node->right;
    while (victim->left) {
        victim = victim->left;
    }
    
    // 分离后继节点
    AVLNode *root = avl_del_easy(victim);
    
    // 与后继节点交换
    *victim = *node;    // left, right, parent
    if (victim->left) {
        victim->left->parent = victim;
    }
    if (victim->right) {
        victim->right->parent = victim;
    }
    
    // 将后继节点附加到父节点，或更新根指针
    AVLNode **from = &root;
    AVLNode *parent = node->parent;
    if (parent) {
        from = parent->left == node ? &parent->left : &parent->right;
    }
    *from = victim;
    return root;
}

AVLNode *avl_offset(AVLNode *node, int64_t offset) {
    int64_t pos = 0;    // 从起始节点的排名差
    while (offset != pos) {
        if (pos < offset && pos + avl_cnt(node->right) >= offset) {
            // 目标在右子树内
            node = node->right;
            pos += avl_cnt(node->left) + 1;
        } else if (pos > offset && pos - avl_cnt(node->left) <= offset) {
            // 目标在左子树内
            node = node->left;
            pos -= avl_cnt(node->right) + 1;
        } else {
            // 转到父节点
            AVLNode *parent = node->parent;
            if (!parent) {
                return NULL;
            }
            if (parent->right == node) {
                pos -= avl_cnt(node->left) + 1;
            } else {
                pos += avl_cnt(node->right) + 1;
            }
            node = parent;
        }
    }
    return node;
} 